from sklearn.preprocessing import OneHotEncoder
import pandas as pd

def data_preprocessing(path, encoders=None, is_train=True):
    """
    数据预处理函数，支持训练集和测试集列顺序一致
    """
    data = pd.read_csv(path)

    # 映射编码 Gender（固定映射）
    data["Gender"] = data["Gender"].map({"Female": 0, "Male": 1})

    # LabelEncoder Attrition（目标变量）
    from sklearn.preprocessing import LabelEncoder
    le = LabelEncoder()
    data['Attrition'] = le.fit_transform(data['Attrition'])

    # One-Hot 编码列
    cat_cols = ['BusinessTravel', 'Department', 'EducationField', 'JobRole', 'MaritalStatus', 'OverTime']

    if is_train:
        # 训练模式：拟合 One-Hot Encoder
        ohe = OneHotEncoder()
        encoded = ohe.fit_transform(data[cat_cols])
        encoded_df = pd.DataFrame(encoded.toarray(), columns=ohe.get_feature_names_out(cat_cols))
        data = pd.concat([data.reset_index(drop=True), encoded_df], axis=1)
        data = data.drop(cat_cols, axis=1)

        # 保存编码器和原始列结构
        encoders = {
            'ohe': ohe,
            'cat_cols': cat_cols,
            'all_columns': data.columns.tolist()
        }
    else:
        # 测试模式：使用训练集编码器
        ohe = encoders['ohe']
        encoded = ohe.transform(data[cat_cols])
        encoded_df = pd.DataFrame(encoded.toarray(), columns=ohe.get_feature_names_out(cat_cols))
        data = pd.concat([data.reset_index(drop=True), encoded_df], axis=1)
        data = data.drop(cat_cols, axis=1)

        # 强制使用训练集的列顺序
        all_columns = encoders['all_columns']
        data = data.reindex(columns=all_columns, fill_value=0)  # 补0缺失列

    return data, encoders